{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a93dbafc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# float 型 tensor 转换为 list，保留小数点后 decimals 位小数 的演示\n",
    "# tensor.tolist() 生成的 list 实际上会输出所有小数位，且值不对\n",
    "import json\n",
    "import torch\n",
    "import numpy as np\n",
    "arr = torch.tensor([[1.123456, 2.1234567, 3.12345678], [21.1232323, 22.2324234234, 23.234234234]], device=\"cuda\")\n",
    "\n",
    "\n",
    "def ftensor2list(tensor: torch.Tensor, decimals: int = 4) -> list:\n",
    "    \"\"\"\n",
    "    float 型 tensor 转换为 list，保留小数点后 decimals 位小数\n",
    "    Args:\n",
    "        tensor:     float 型 tensor\n",
    "        decimals:   保留小数点后 decimals 位小数\n",
    "    Returns:\n",
    "        list\n",
    "    \"\"\"\n",
    "    l = tensor.tolist()\n",
    "    arr = np.array(l)\n",
    "    # arr = torch.tensor(l)  # 不能使用 torch.tensor(confs) 会导致输出所有小数位，且值不对\n",
    "    ll = arr.round(decimals=decimals).tolist()\n",
    "    return ll\n",
    "\n",
    "l = ftensor2list(arr)\n",
    "print(f\"list = {l}\")\n",
    "print(\"=\"* 40)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9302fbd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 保存 kr 格式, 不能保存 masks，masks 太大，保存下来需要 90MB+；无法打开解析\n",
    "def ftensor2list(tensor: torch.Tensor, decimals: int = 5) -> list:\n",
    "    \"\"\"\n",
    "    float 型 tensor 转换为 list，保留小数点后 decimals 位小数\n",
    "    Args:\n",
    "        tensor:     float 型 tensor\n",
    "        decimals:   保留小数点后 decimals 位小数\n",
    "    Returns:\n",
    "        list\n",
    "    \"\"\"\n",
    "    l = tensor.tolist()\n",
    "    arr = np.array(l)\n",
    "    # arr = torch.tensor(l)  # 不能使用 torch.tensor(confs) 会导致输出所有小数位，且值不对\n",
    "    ll = arr.round(decimals).tolist()\n",
    "    return ll\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "print(f\"保存 kr 文件: {kr_file}\")\n",
    "decimals = 5\n",
    "kr = results[0]\n",
    "classes = kr.boxes.cls.tolist()\n",
    "confs = ftensor2list(kr.boxes.conf, decimals)  # (torch.round(kr.boxes.conf.to(\"cpu\") * 1e5)/1e5).tolist()\n",
    "boxes = ftensor2list(kr.boxes.data, decimals)  # ((kr.boxes.data.to(\"cpu\") * decimals).round() / decimals).tolist()\n",
    "segments =  kr.masks.xy\n",
    "masks = ftensor2list(kr.masks.data, decimals)  # kr.masks.data.tolist()\n",
    "# keypoints = kr.keypoints.data.tolist()\n",
    "# obb = kr.obb.data.tolist()\n",
    "print(f\"confs = {confs}\")\n",
    "\n",
    "# print(segments[0].shape)\n",
    "\n",
    "print(f\"class num = {len(classes)}\")\n",
    "print(f\"conf num = {len(confs)}\")\n",
    "print(f\"box num = {len(boxes)}\")\n",
    "print(f\"segment num = {len(segments)}\")\n",
    "print(f\"mask num = {len(masks)}\")\n",
    "\n",
    "def get_box(box: list) -> dict:\n",
    "    bb: dict={\"x1\": box[0], \"y1\": box[1], \"x2\": box[2], \"y2\": box[3]}\n",
    "    return bb\n",
    "\n",
    "kr_dict: list = []\n",
    "for (cls_id, conf, box, segment, mask)  in zip(classes, confs, boxes, segments, masks):\n",
    "    name = kr.names[cls_id]\n",
    "    segment_x = ftensor2list(segment[:, 0], decimals) \n",
    "    segment_y = ftensor2list(segment[:, 1], decimals) \n",
    "    # print(f\"segment len = {len(segment)}\")\n",
    "    kr_dict.append({\n",
    "        \"name\": name,\n",
    "        \"class\": int(cls_id),\n",
    "        \"confidence\": conf,\n",
    "        \"box\": get_box(box),\n",
    "        \"segments\": {\"x\": segment_x, \"y\": segment_y}, #  segment,\n",
    "        \"masks\": mask,\n",
    "    })\n",
    "\n",
    "kr_dict\n",
    "\n",
    "# 不能保存 masks，masks 太大，保存下来需要 90MB+；无法打开解析\n",
    "# kr_json = json.dumps(kr_dict, indent=4, ensure_ascii=False)\n",
    "# kr_json = json.dumps(kr_dict, ensure_ascii=False)\n",
    "# with open(kr_file, \"w\", encoding=\"utf-8\") as f:\n",
    "#     f.write(kr_json)\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
